from __future__ import annotations
from typing import Sequence
import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
import manim as mnm

from investment.date_utilities import get_date_by_year, format_month_interval, DateUtility
from investment.operation import calc_reward_rate
from investment.time_series import IndexTimeSeries
from investment.constants import A_INDEX_DICT
from visual.common_scenes import MY_FONT, TimeSeriesAxes
from quantify.find_diamond import DiamondMatrix
from misc.character_design import *
from myMovie_emotion import EmotionSeriesSegments


def max_min_norm(value, v_min, v_max):
    assert (v_max - v_min) > 1e-4
    return (value - v_min) / (v_max - v_min)


class TimeSeriesTwinAxes(TimeSeriesAxes):
    def __init__(self, value_series: np.ndarray, title, height, year=False):
        super(TimeSeriesTwinAxes, self).__init__(value_series, title, height, year)
        self.twin_axis: mnm.NumberLine = None
        self.twin_series = None
        self.twin_axis_step = 10
        self.twin_dot = None
        self.twin_trace = None
        self.twin_wandering = None

    def set_twin_axis(self, twin_value_series: np.ndarray):
        self.twin_series = twin_value_series
        self.twin_axis = mnm.NumberLine(
            x_range=[0, twin_value_series.max(), self.twin_axis_step],
            length=self.axes.get_axis(1).length,
            include_numbers=True,
            font_size=self.axes.get_axis(1).font_size,
            label_direction=mnm.RIGHT,
            rotation=self.axes.get_axis(1).rotation,
            color=mnm.BLUE
        )
        x_axis = self.axes.get_axis(0)
        self.twin_axis.shift(x_axis.n2p(x_axis.x_range[1]) - self.twin_axis.number_to_point(0))
        self.axes_group.add(self.twin_axis)

    def coords_to_point_twinx(self, *coords: Sequence[float] | np.ndarray):
        assert self.twin_axis is not None, RuntimeError("have not set twin axis")
        ax0 = self.axes.get_x_axis()
        origin = ax0.number_to_point(ax0.x_min)
        coords = np.asarray(coords)
        result = np.array(origin)
        for axis, number in zip([ax0, self.twin_axis], coords):
            result += axis.number_to_point(number) - axis.number_to_point(axis.x_min)
        return result

    def index_to_point_twinx(self, index):
        return self.coords_to_point_twinx(index, self.twin_series[index])

    def get_twin_series(self, index):
        return self.twin_series[index]

    def wandering_series_twin(self, dot_color=mnm.BLUE) -> tuple[mnm.VGroup, mnm.Animation]:
        # cp_color = mnm.GREY
        zero_point = self.index_to_point_twinx(0)
        self.twin_dot = mnm.Dot(zero_point, color=dot_color)
        self.twin_trace = mnm.TracedPath(self.twin_dot.get_center, stroke_color=dot_color)
        self.twin_wandering = mnm.VMobject(stroke_color=dot_color, stroke_width=1)
        self.twin_wandering.set_points_as_corners([zero_point, zero_point])
        for idx in range(1, self.sample_length):
            self.twin_wandering.add_line_to(self.index_to_point_twinx(idx))
        self.axes_group.add(self.twin_dot, self.twin_trace)
        return self.axes_group, mnm.MoveAlongPath(self.twin_dot, self.twin_wandering)

    def clear_dot_twin(self):
        self.axes_group -= self.twin_dot
        self.axes_group -= self.twin_trace
        self.axes_group += self.twin_wandering

    def get_evaluation_line(self, e_value):
        st = self.coords_to_point_twinx(0, e_value)
        ed = self.coords_to_point_twinx(self.axes.x_range[1], e_value)
        dash_line = mnm.DashedLine(st, ed, dashed_ratio=0.2, color=mnm.BLUE)
        return dash_line


class ColorBar(mnm.Group):
    def __init__(self, right_bound, width=0.3, height=mnm.config.frame_height * 0.8):
        super().__init__()
        assert width < height
        cmap = plt.get_cmap("jet")
        gradient = np.linspace(1, 0, 256)
        ratio = width / height
        col = int(256 * ratio)
        gradient = np.dot(np.resize(gradient, (256, 1)), np.resize([1] * col, (1, col)))
        self.im = mnm.ImageMobject(filename_or_array=(cmap(gradient) * 255).astype(np.uint8))
        self.im.height = height
        self.im.next_to(right_bound, direction=mnm.RIGHT)
        self.axis = None
        self.add(self.im)

    def set_axis(self, value_range: Sequence):
        value_range = np.round(value_range, 1)
        self.axis = mnm.NumberLine(
            x_range=value_range,
            length=self.im.height,
            include_numbers=True,
            tick_size=0.05,
            longer_tick_multiple=1,
            font_size=16,
            label_direction=mnm.RIGHT,
            line_to_number_buff=mnm.SMALL_BUFF,
            rotation=mnm.PI / 2,
            color=mnm.WHITE
        ).move_to(self.im.get_center())
        self.axis.shift(mnm.RIGHT * (self.im.width + self.axis.width) / 2.)
        self.add(self.axis)
        return self.axis

    def get_axis_value(self, value):
        assert isinstance(self.axis, mnm.NumberLine)
        return self.axis.number_to_point(value)


def get_num_of_range(x):
    return int((x[1] - x[0]) / x[2] + 1)


class HeatImage(mnm.Group):
    def __init__(self, x_range, y_range, height, x_labels: dict, y_labels: dict, **kwargs):
        super().__init__()
        self.plane = mnm.NumberPlane(x_range=x_range, y_range=y_range,
                                     x_length=height, y_length=height,
                                     x_axis_config={"label_direction": mnm.DOWN},
                                     y_axis_config={"include_numbers": False, "label_direction": mnm.LEFT,
                                                    "font_size": 16},
                                     background_line_style={"stroke_color": mnm.GRAY, "stroke_opacity": 0.2},
                                     **kwargs
                                     )
        # mnm.ImageMobject.__init__(self, filename_or_array=arr, **kwargs)
        self.height = height
        self.grid_width = height / get_num_of_range(self.x_range)
        self.grid_height = height / get_num_of_range(self.y_range)
        # self.set_resampling_algorithm(mnm.RESAMPLING_ALGORITHMS["box"])
        self.cmap = plt.get_cmap("jet")
        self.colorbar = ColorBar(self.plane.get_center() + mnm.RIGHT * self.plane.width / 2, height=self.y_length)
        self.value_scaler = []
        self.x_axis.add_labels(x_labels, font_size=16, buff=mnm.SMALL_BUFF/2)
        self.y_axis.add_labels(y_labels, font_size=16)
        self.add(self.plane, self.colorbar)

    @mnm.override_animation(mnm.Create)
    def _create_override(self):
        # 装饰器模式, getattr callable function of object
        anim = mnm.AnimationGroup(mnm.Create(self.x_axis), mnm.Create(self.y_axis),
                                  mnm.Create(self.background_lines), mnm.FadeIn(self.colorbar))
        return anim

    def mapping_color(self, value):
        assert len(self.value_scaler) > 0
        value = max_min_norm(value, self.value_scaler[0], self.value_scaler[1])
        color = self.cmap(value)
        color = mpl.colors.to_hex(color)
        return color

    def get_colorbar_point(self, value):
        return self.colorbar.get_axis_value(value)

    def coord_to_point(self, coord):
        return self.plane.coords_to_point(coord[0], coord[1])

    def set_color_range(self, color_range):
        """
        colorbar value max and min
        :param color_range: [min, max, step]
        :return:
        """
        self.colorbar.set_axis(color_range)
        self.value_scaler = color_range

    def __getattr__(self, item):
        # 适配器，适配所有NumberPlane 成员
        if item in self.plane.__dict__:
            return getattr(self.plane, item)
        else:
            raise AttributeError('Unknown attribute : %s' % item)


class HeatTest(mnm.Scene):
    def construct(self):
        with open(r"D:\Work\Projects\Jazz\test\data\mat_amax.npy", 'rb') as f:
            mat = np.load(f)
            mat = mat.T
        img = HeatImage([0, mat.shape[0], 1], [0, mat.shape[1], 1], 5, {1: "1y", 10: "2y"})
        img.set_color_range((np.nanmin(mat), np.nanmax(mat), (np.nanmax(mat) - np.nanmin(mat)) / 20))
        self.play(mnm.Create(img))
        # for i in range(mat.shape[0]):
        #     for j in range(mat.shape[1]):
        #         self.add(img.draw_pixel(i+1, j+1, mat[i, j]))
        self.wait(1)


class DiamondImage:
    def __init__(self, diamond_mat):
        assert isinstance(diamond_mat, DiamondMatrix)
        evaluation_spaces = diamond_mat.evaluation_space
        month_count = diamond_mat.shape[1]
        date_ticks = {year: f"{int(year / 12)}yr" for year in range(0, month_count, 12)}
        pe_tick_step = 2
        pe_ticks = {i * pe_tick_step: e for i, e in enumerate(np.round(evaluation_spaces[::pe_tick_step], 1))}

        self.diamond_mat = diamond_mat
        self.heat_img = HeatImage([0, month_count, 1], [0, diamond_mat.shape[0], 1],
                                  mnm.config.frame_height * 0.5, date_ticks, pe_ticks)
        self.color_ticks_number = 10
        self.max_matrix = None
        self.min_matrix = None
        self.pixels_groups = {}

    def shift(self, vec):
        self.heat_img.shift(vec)

    def set_heat_color(self, matrix):
        self.heat_img.set_color_range(
            (np.nanmin(matrix), np.nanmax(matrix), (np.nanmax(matrix) - np.nanmin(matrix)) / self.color_ticks_number))

    def set_heat_color_mode(self, mode):
        if mode == "max":
            mat = self._max()
        elif mode == "min":
            mat = self._min()
        else:
            raise RuntimeError("error mode")
        self.heat_img.set_color_range(
            (np.nanmin(mat), np.nanmax(mat), (np.nanmax(mat) - np.nanmin(mat)) / self.color_ticks_number))

    def remove_all_pixels(self):
        kys = list(self.pixels_groups.keys())
        for ky in kys:
            self.pixels_groups.pop(ky)

    def clean_pixels_from_scene(self, screen):
        kys = list(self.pixels_groups.keys())
        for ky in kys:
            screen.remove(self.pixels_groups[ky])

    def animate_draw_operation(self, operation):
        assert isinstance(operation, Operation)
        color = self.heat_img.mapping_color(operation.reward_rate)
        operation.rect.color = color
        if operation.reward_text:
            creation_anim = []
            framebox = mnm.SurroundingRectangle(operation.reward_text)
            creation_anim.append(mnm.ShowPassingFlash(framebox))
            color_pt = self.heat_img.get_colorbar_point(operation.reward_rate)
            creation_anim.append(
                mnm.Create(mnm.Arrow(color_pt, operation.rect.get_corner(mnm.DL), stroke_width=3),
                           remover=True))
            creation_anim.append(mnm.FadeIn(operation.rect))
            creation_anim = mnm.AnimationGroup(*creation_anim)
        else:
            creation_anim = operation.mobs_group.animate.become(operation.rect)
        return creation_anim

    def actual_value_to_point(self, actual_values):
        """
        convert actualvalues(month, evaluation) to pixel coord by diamond image. and return point
        :param actual_values: coordinates
        :param diamond_img:
        :return: point in scene. np.ndarray
        """
        assert len(actual_values) == 2
        return self.heat_img.coord_to_point(self.actual_value_to_coord(actual_values[0], actual_values[1]))

    def actual_value_to_coord(self, elapsed_month, evaluation):
        return tuple([elapsed_month - 1, self.diamond_mat.which_bins(evaluation)])

    def animate_move_to_pixel(self, operation):
        # month range start from 1~, but coord start from 0, so month2coord need -1
        coord = self.actual_value_to_coord(operation.elapsed_month, operation.start_sample.evaluation)
        moving_anim = []
        if operation.evaluation_text:
            text_anim = mnm.AnimationGroup(
                mnm.ShowPassingFlash(mnm.SurroundingRectangle(operation.evaluation_text)),
                mnm.ShowPassingFlash(mnm.SurroundingRectangle(operation.elapsed_text)))
            moving_anim.append(text_anim)

        rect = operation.rect.copy()
        color = self.heat_img.mapping_color(operation.reward_rate)
        self.pixels_groups[coord] = mnm.Rectangle(color, width=self.grid_width,
                                                  height=self.grid_height, fill_opacity=1)
        # note: align to UP RIGHT corner.
        self.pixels_groups[coord].align_to(self.heat_img.coord_to_point(coord), direction=mnm.DL)
        moving_anim.append(rect.animate.become(self.pixels_groups[coord]))
        moving_anim = mnm.Succession(*moving_anim)
        self.pixels_groups[coord] = rect
        return moving_anim

    def __getattr__(self, item):
        # 适配器，适配所有HeatImg 成员
        if item in self.heat_img.__dict__:
            return getattr(self.heat_img, item)
        else:
            raise AttributeError('Unknown attribute : %s' % item)

    def _max(self):
        if self.max_matrix is None:
            self.max_matrix = self.diamond_mat.get_statistic_matrix(func=np.max)
        return self.max_matrix

    def _min(self):
        if self.min_matrix is None:
            self.min_matrix = self.diamond_mat.get_statistic_matrix(func=np.min)
        return self.min_matrix

    def show_mat(self, mode="max"):
        matrix = self.max_matrix if mode == "max" else self.min_matrix
        assert matrix is not None
        self.remove_all_pixels()
        for i in range(matrix.shape[0]):
            for j in range(matrix.shape[1]):
                value = matrix[i, j]
                if value is np.nan:
                    continue
                # note pe is y_axis, date is x_axis. but shape[0] is pe, shape[1] is time. so transpose.
                coord = (j, i)
                self.pixels_groups[coord] = mnm.Rectangle(self.heat_img.mapping_color(value),
                                                          width=self.grid_width,
                                                          height=self.grid_height, fill_opacity=1)
                self.pixels_groups[coord].align_to(self.heat_img.coord_to_point(coord), direction=mnm.DL)
                self.heat_img.add(self.pixels_groups[coord])


class SampleOfSeries:
    def __init__(self, time_axes, index):
        assert isinstance(time_axes, TimeSeriesTwinAxes)
        self.index = index
        self.price = time_axes.get_series(index)
        self.evaluation = time_axes.get_twin_series(index)
        self.price_pt = time_axes.index_to_point(index)
        self.evaluation_pt = time_axes.index_to_point_twinx(index)


def check_interval(date_utility, start_index, month_interval: int):
    end_idx = date_utility.get_interval_trading_index(start_index, month_interval)
    return end_idx is not None


class Operation:
    def __init__(self, time_axes, date_utility, start_index, month_interval: int):
        assert isinstance(time_axes, TimeSeriesTwinAxes)
        assert isinstance(date_utility, DateUtility)
        self.date_utility = date_utility
        self.start_sample = SampleOfSeries(time_axes, start_index)
        self.start_day = self.date_utility.trading_dates[start_index]

        self.end_index = self.date_utility.get_interval_trading_index(start_index, month_interval)
        if self.end_index is None:
            raise ValueError("end index out of range by time interval.")
        self.end_day = self.date_utility.get_trading_day(self.end_index)
        self.elapsed_month = month_interval
        self.end_sample = SampleOfSeries(time_axes, self.end_index)
        self.reward_rate = calc_reward_rate(self.start_sample.price, self.end_sample.price)

        # manim mobjects
        radius = mnm.DEFAULT_SMALL_DOT_RADIUS
        self.st_price_dot = mnm.Dot(self.start_sample.price_pt, radius=radius)
        self.st_evaluation_dot = mnm.Dot(self.start_sample.evaluation_pt, radius=radius)
        self.st_p2e_line = mnm.DashedLine(self.start_sample.price_pt, self.start_sample.evaluation_pt)
        self.ed_price_dot = mnm.Dot(self.end_sample.price_pt, radius)
        self.elapsed_time_line = mnm.DashedLine(self.start_sample.price_pt, np.array(
            [self.end_sample.price_pt[0], self.start_sample.price_pt[1], 0.]))
        self.reward_rate_line = mnm.DashedLine(self.end_sample.price_pt, np.array(
            [self.end_sample.price_pt[0], self.start_sample.price_pt[1], 0.]))
        self.mobs_group = mnm.VGroup(self.st_price_dot, self.st_evaluation_dot, self.st_p2e_line, self.ed_price_dot,
                                     self.elapsed_time_line, self.reward_rate_line)
        # manim text
        self.evaluation_text = None
        self.reward_text = None
        self.elapsed_text = None

        # rect mob
        self.rect = None

    def animation_creation(self, text_flag=False, run_time=1):
        """
        Animation groups of pipeline. from one sample/interval of operation mapping to heat img.
        :param run_time:
        :param text_flag:
        :return: list of animations
        """
        eval_animation = [mnm.FadeIn(self.st_price_dot, self.st_p2e_line, self.st_evaluation_dot)]
        if text_flag:
            self.evaluation_text = mnm.Text(f"市盈率:{self.start_sample.evaluation:.0f} (估值)", font_size=18,
                                            font=MY_FONT).next_to(self.st_evaluation_dot)
            self.mobs_group.add(self.evaluation_text)
            eval_animation.append(mnm.FadeIn(self.evaluation_text))
        eval_animation = mnm.AnimationGroup(*eval_animation)
        elapsed_animation = [mnm.FadeIn(self.elapsed_time_line, self.reward_rate_line, self.ed_price_dot)]
        if text_flag:
            self.elapsed_text = mnm.Text(f"{format_month_interval(self.elapsed_month)} (持有时间)",
                                         font_size=18, font=MY_FONT).next_to(self.elapsed_time_line, direction=mnm.UP)
            elapsed_animation.append(mnm.FadeIn(self.elapsed_text))
            self.reward_text = mnm.Text(f"{self.reward_rate:.1%} (回报率)", font_size=18, font=MY_FONT).next_to(
                self.reward_rate_line)
            self.mobs_group.add(self.elapsed_text)
            self.mobs_group.add(self.reward_text)
            elapsed_animation.append(mnm.FadeIn(self.reward_text))
        elapsed_animation = mnm.AnimationGroup(*elapsed_animation)
        return mnm.Succession(eval_animation, elapsed_animation, run_time=run_time)

    def clean_up_from_scene(self, screen):
        screen.remove(*self.mobs_group.submobjects)
        screen.remove(self.rect)


class VRectGroup:
    def __init__(self):
        self.col_num = 4
        self.row_num = 8
        width = 2
        height = 4
        frame_position = 2.5 * mnm.DOWN + 0.3 * mnm.RIGHT
        self.grid_xstep = width / self.col_num
        self.grid_ystep = height/self.row_num
        self.frame = mnm.Rectangle(width=width, height=height,
                                   grid_xstep=self.grid_xstep, grid_ystep=self.grid_ystep).move_to(frame_position)
        self.rect_group = mnm.VGroup()
        self.value_list = []
        self.text_group = mnm.VGroup()
        self.max_text = mnm.Text("最大", font=MY_FONT, font_size=12)
        self.min_text = mnm.Text("最小", font=MY_FONT, font_size=12)

    def create_rects(self, operation_list):
        num = len(operation_list)
        if num == 0:
            return
        for i in range(num):
            self.value_list.append(operation_list[i].reward_rate)
            operation_list[i].rect = mnm.Rectangle(width=self.grid_xstep/2, height=self.grid_ystep/2, fill_opacity=1)
            self.rect_group.add(operation_list[i].rect)
        self.rect_group.arrange_in_grid(cols=self.col_num, buff=mnm.MED_SMALL_BUFF)
        self.rect_group.align_to(self.frame, direction=mnm.UL)
        for i in range(num):
            self.text_group.add(mnm.Text(f"{operation_list[i].reward_rate:.1%}", font=MY_FONT, font_size=12).move_to(
                operation_list[i].rect))

    def get_maximum_animate(self) -> tuple[int, mnm.Animation]:
        if len(self.value_list) == 0:
            return -1, mnm.Wait(0.1)
        max_idx = np.nanargmax(self.value_list)
        max_rect = self.rect_group.submobjects[max_idx]
        self.max_text.next_to(max_rect, direction=mnm.DOWN, buff=mnm.SMALL_BUFF)
        return max_idx, mnm.AnimationGroup(mnm.FadeIn(self.max_text), mnm.FadeIn(self.text_group))

    def get_minimum_animate(self) -> tuple[int, mnm.Animation]:
        if len(self.value_list) == 0:
            return -1, mnm.Wait(0.1)
        min_idx = np.nanargmin(self.value_list)
        min_rect = self.rect_group.submobjects[min_idx]
        self.min_text.next_to(min_rect, direction=mnm.DOWN, buff=mnm.SMALL_BUFF)
        return min_idx, mnm.FadeIn(self.min_text)

    def clean_up_from_scene(self, screen):
        screen.remove(self.rect_group)
        screen.remove(self.max_text)
        screen.remove(self.min_text)
        screen.remove(self.text_group)
        self.rect_group.remove(*self.rect_group.submobjects)
        self.text_group.remove(*self.text_group.submobjects)
        self.value_list = []


class AnnotationGroup(mnm.Group):
    def __init__(self, font_size=18):
        self.font_size = font_size
        txt1 = mnm.Text("1. 高估值买入很难赚钱", font_size=font_size, font=MY_FONT)
        txt2 = mnm.Text("2. 低估值买入有择时的优势,可能有更高回报", font_size=font_size, font=MY_FONT)
        txt3 = mnm.Text("3. 低估值买入,从最小稳定收益来说,时间是朋友", font_size=font_size, font=MY_FONT)
        txt4 = mnm.Text("4. 极端低估的时间太少,导致取不到样本. 相信是块钻石", font_size=font_size, font=MY_FONT)
        super().__init__(txt1, txt2, txt3, txt4)
        self.arrange(direction=mnm.DOWN, buff=mnm.MED_LARGE_BUFF)
        self.txt_list = list(txt.align_on_border(direction=mnm.LEFT, buff=mnm.LARGE_BUFF)
                             for txt in [txt1, txt2, txt3, txt4])

        self.annotation_area = [
            {"max": [(1, 38), (100, 50)], "min": [(1, 38), (100, 50)]},
            {"max": [(12, 12), (119, 20)], "min": None},
            {"max": None, "min": [(60, 9), (110, 15)]},
            {"max": [(108, 8.1), (119, 9.5)], "min": [(108, 8.1), (119, 9.5)]},
        ]

    def __iter__(self):
        return iter(self.txt_list)


def create_annotation_circle(dl: np.ndarray, ur: np.ndarray, txt: mnm.Text):
    circle_center = (dl + ur)/2
    width, height, _ = np.abs(dl-ur)
    circle = mnm.Ellipse(width+mnm.SMALL_BUFF, height+mnm.SMALL_BUFF).move_to(circle_center)
    end_pt = txt.get_corner(direction=mnm.UR)
    direction = end_pt - circle.get_center()
    circle_start = circle.get_boundary_point(direction=direction)
    arrow = mnm.Arrow(circle_start, end_pt)
    animation = mnm.AnimationGroup(mnm.Create(circle), mnm.Create(arrow))
    return animation


class AnnotationGroupTest(mnm.Scene):
    def construct(self):
        group = AnnotationGroup()
        for txt in iter(group):
            txt = txt
            self.play(mnm.Succession(create_annotation_circle(mnm.UP*2, mnm.UR*3, txt), mnm.FadeIn(txt)))


class MyMovie(mnm.Scene):
    def __init__(self, time_series, last_time=2, axes_height_rate=0.4, axes_width_rate=1):
        assert isinstance(time_series, IndexTimeSeries)
        self.time_series = time_series
        self.axes_width = mnm.config.frame_width * axes_width_rate
        self.axes_height = mnm.config.frame_height * axes_height_rate
        self.last_time = last_time
        self.time_axes: mnm.Axes
        self.dates_labels = {}
        # config frame_rate maybe interpolate line details
        self.debug = False
        self.date_utility = DateUtility(self.time_series.dates_series)
        super().__init__()

    def construct(self):
        title = mnm.Text("寻找沪深300的钻石:", font_size=24, font=MY_FONT)
        title.to_edge(mnm.UL, buff=mnm.MED_LARGE_BUFF)
        self.play(mnm.Write(title))
        self.show_prices()
        self.next_section("show pe value")
        self.show_pe()
        self.next_section("create diamond table")
        self.map_diamond()
        self.next_section("show diamond samples")
        self.show_diamond_samples()
        self.next_section("reference")
        self.reference()
        self.wait(1)

    def show_prices(self):
        # add axes
        cp = self.time_series.get_fundamental_metrics("CP")
        title_name = A_INDEX_DICT.inv[self.time_series.stock_code]
        self.time_axes = TimeSeriesTwinAxes(cp, title_name, self.axes_height, year=True)
        year_indexes, year_date = get_date_by_year(self.time_series.dates_series)
        self.time_axes.add_date_ticks(dict(zip(year_indexes, year_date)))
        self.add(self.time_axes.axes_group)
        # show cp
        if self.debug:
            return
        mob_group, wandering_animate = self.time_axes.wandering_series()
        legend_txt = mnm.Text("指数", color=self.time_axes.dot.color,
                              font_size=18, font=MY_FONT).next_to(self.time_axes.dot)
        legend_txt.add_updater(lambda x: x.next_to(self.time_axes.dot))
        self.add(mob_group, legend_txt)
        self.play(wandering_animate, run_time=self.last_time)
        self.time_axes.clear_dot()

    def show_pe(self):
        # add pe axes
        pe = self.time_series.get_fundamental_metrics("PE")
        self.time_axes.set_twin_axis(pe)
        self.play(mnm.FadeIn(self.time_axes.twin_axis))
        # show pe
        if self.debug:
            return
        _, wandering_animate = self.time_axes.wandering_series_twin()
        legend_txt = mnm.Text("市盈率", color=self.time_axes.twin_dot.color,
                              font_size=16, font=MY_FONT).next_to(self.time_axes.twin_dot)
        legend_txt.add_updater(lambda x: x.next_to(self.time_axes.twin_dot))
        self.add(legend_txt)
        self.play(wandering_animate, run_time=self.last_time)
        self.time_axes.clear_dot_twin()

    def map_diamond(self):
        self.clear()
        up_length = self.time_axes.axes_group.get_boundary_point(mnm.DOWN)[1]
        self.play(self.time_axes.axes_group.animate.to_edge(mnm.UP, mnm.MED_SMALL_BUFF))
        gold_splitter = 0.7
        counts_of_all = len(self.time_series.dates_series)
        sample_index = int(counts_of_all * (1 - gold_splitter))
        operation = Operation(self.time_axes, self.date_utility, sample_index, month_interval=6 * 12)
        self.play(operation.animation_creation(True))

        # show table
        if True:
            import pickle
            with open(r"D:\Work\Projects\Jazz\test\data\diamond_mat.txt", 'rb') as f:
                diamond_mat = pickle.load(f)
        else:
            pass
            cp = self.time_series.get_fundamental_metrics("CP")
            pe = self.time_series.get_fundamental_metrics("PE")
            diamond_mat = DiamondMatrix(pe, cp, self.date_utility)
        diamond_img = DiamondImage(diamond_mat)
        diamond_img.shift([-4., up_length, 0.])
        diamond_img.set_heat_color_mode("max")
        self.play(mnm.Create(diamond_img.heat_img))
        vrect_group = VRectGroup()
        vrect_group.create_rects([operation])
        draw_operation_anim = diamond_img.animate_draw_operation(operation)
        self.play(draw_operation_anim)
        self.play(diamond_img.animate_move_to_pixel(operation))
        vrect_group.clean_up_from_scene(self)
        operation.clean_up_from_scene(self)
        draw_operation_anim.clean_up_from_scene(self)
        self.wait(0.5)

        diamond_img_min = DiamondImage(diamond_mat)
        diamond_img_min.shift([4, up_length, 0.])
        diamond_img_min.set_heat_color_mode("min")
        self.add(diamond_img_min.heat_img)
        pe_loop = [20, 40]
        month_loop = [12*3, 12*6, 12*9]
        loop_count = 0
        for pe in pe_loop:
            pe_line = self.time_axes.get_evaluation_line(pe)
            self.play(mnm.Create(pe_line))  # show pe price line
            for month in month_loop:
                pe_idx_list = diamond_mat.get_evaluation_bucket(pe)
                ope_list = []
                for idx in pe_idx_list:
                    if check_interval(self.date_utility, idx, month):
                        ope_list.append(Operation(self.time_axes, self.date_utility, idx, month))
                if len(ope_list) == 0:
                    continue
                if (loop_count + 1) % len(month_loop) == 1:
                    highlight = {str(pe): mnm.YELLOW, str(month//12): mnm.YELLOW}
                else:
                    highlight = {str(pe): mnm.WHITE, str(month//12): mnm.YELLOW}
                highwight = {str(pe): mnm.BOLD, str(month // 12): mnm.BOLD}
                txt = mnm.Text(f"市盈率: {pe}, 持有时间: {month//12} 年", t2c=highlight, t2w=highwight, disable_ligatures=True,
                               font_size=16, font=MY_FONT).next_to(vrect_group.frame, direction=mnm.UP, buff=0.2).\
                    align_to(vrect_group.frame, direction=mnm.LEFT)
                self.play(mnm.FadeIn(txt))
                vrect_group.create_rects(ope_list)
                operation_anims = mnm.AnimationGroup(*list(map(lambda ope: ope.animation_creation(), ope_list)), remover=True)
                self.play(operation_anims)
                draw_pixel_anims = mnm.AnimationGroup(*list(map(diamond_img.animate_draw_operation, ope_list)))
                self.play(draw_pixel_anims)
                max_idx, text_anim = vrect_group.get_maximum_animate()
                self.play(text_anim)
                self.play(diamond_img.animate_move_to_pixel(ope_list[max_idx]))
                min_idx, text_anim = vrect_group.get_minimum_animate()
                self.play(text_anim)
                self.play(diamond_img_min.animate_move_to_pixel(ope_list[min_idx]))
                vrect_group.clean_up_from_scene(self)
                list(map(lambda ope: ope.clean_up_from_scene(self), ope_list))
                operation_anims.clean_up_from_scene(self)
                self.remove(txt)
                loop_count += 1
            self.remove(pe_line)

        # note: don't add diamond_img by mistake, which will cause error in self.wait(), uneasy to debug.
        self.wait(0.5)
        self.add(diamond_img.heat_img)
        diamond_img.clean_pixels_from_scene(self)
        diamond_img_min.clean_pixels_from_scene(self)
        diamond_img.show_mat("max")
        diamond_img_min.show_mat("min")
        self.wait(0.5)

        # annotate
        shift_up_length = -2*up_length+mnm.MED_LARGE_BUFF
        self.play(self.time_axes.axes_group.animate.shift([0, shift_up_length, 0]),
                  diamond_img.heat_img.animate.shift([0, shift_up_length, 0]),
                  diamond_img_min.heat_img.animate.shift([0, shift_up_length, 0]))
        group = AnnotationGroup().next_to(diamond_img.heat_img, direction=mnm.DOWN, buff=mnm.MED_LARGE_BUFF)
        for i, txt in enumerate(group):
            if group.annotation_area[i].get("max"):
                value_dl, value_ur = group.annotation_area[i].get("max")
                self.play(create_annotation_circle(diamond_img.actual_value_to_point(value_dl),
                                                   diamond_img.actual_value_to_point(value_ur), txt))
            if group.annotation_area[i].get("min"):
                value_dl, value_ur = group.annotation_area[i].get("min")
                self.play(create_annotation_circle(diamond_img_min.actual_value_to_point(value_dl),
                                                   diamond_img_min.actual_value_to_point(value_ur), txt))
            self.play(mnm.FadeIn(txt))

    def show_diamond_samples(self):
        self.clear()
        title = mnm.Text("相对低估时(市盈率≈10)买入后:", font_size=24, font=MY_FONT)
        title.to_edge(mnm.UL, buff=mnm.MED_LARGE_BUFF)
        self.add(title)
        emotion_axes = EmotionSeriesSegments(self.time_series, 12)
        for axis in iter(emotion_axes):
            self.add(axis.axes_group, axis.trace, axis.character, axis.redraw_x)
            last_pos = 0
            cache_labels = []
            for key, pos in axis.journey_nodes.items():
                self.play(axis.animate_moving(last_pos, pos))
                self.play(axis.animate_character_emotion(last_pos, pos))
                self.play(axis.remove_bubble())
                cache_pos = axis.cache_position()
                cache_labels.append(cache_pos)
                self.add(cache_pos)
                last_pos = pos
            self.play(mnm.Create(axis.get_duration_text()))
            self.wait(1)
            self.remove(axis.axes_group, axis.trace, axis.dot, axis.character, axis.redraw_x, axis.duration_text, *cache_labels)
        conclusion = mnm.Text("幸福是向山顶攀登的过程中的种种经历和感受. 耐心持有,静待花开.", font_size=24, font=MY_FONT)
        self.play(mnm.Write(conclusion.next_to(title, direction=mnm.DOWN, buff=mnm.LARGE_BUFF).to_edge(
            mnm.LEFT, mnm.MED_LARGE_BUFF)))
        self.wait(1)

    def reference(self):
        self.clear()
        title = mnm.Text("参考", font=MY_FONT).scale(1.5)
        title.to_edge(mnm.UP, buff=mnm.MED_SMALL_BUFF)
        randy, morty = mnm.VGroup(Yan(), Yan(color=mnm.GRAY))
        for pi, vect in (randy, mnm.LEFT), (morty, mnm.RIGHT):
            pi.set(height=title.height)
            # pi.change_mode("thinking")
            pi.look_at(mnm.DOWN)
            pi.next_to(title, vect, buff=mnm.MED_LARGE_BUFF)
        line = mnm.DashedLine(mnm.config.frame_x_radius * mnm.LEFT, mnm.config.frame_x_radius * mnm.RIGHT)
        line.next_to(title, direction=mnm.DOWN, buff=mnm.MED_SMALL_BUFF)
        self.add_foreground_mobjects(title, randy, morty, line)

        width = mnm.config.frame_width * 0.8
        data_group = mnm.VGroup(mnm.Text("1. 数据来源：理杏仁开放平台", font=MY_FONT), mnm.MarkupText(
            f'<span underline="single" underline_color="white">https://www.lixinger.com/open/api/</span>',
            color=mnm.WHITE
        )).arrange(mnm.RIGHT)
        ref_group = mnm.VGroup(mnm.Text("2. 想法来源: 无人知晓E02 & 《寻找皇冠上的钻石》--信达证券-陈嘉禾", font=MY_FONT))
        manim_group = mnm.VGroup(mnm.Text("3. 动画编程库ManimCommunity/数学科普频道3b1b", font=MY_FONT),
                                 Logo().scale(0.5), mnm.ManimBanner().scale(0.5)).arrange(direction=mnm.RIGHT)
        wechat = mnm.Text("4. 潜入用户群收到的信号", font=MY_FONT)
        all = mnm.Group(data_group, ref_group, manim_group).arrange(direction=mnm.DOWN, buff=1.5)
        all.set(width=width)
        self.play(mnm.Write(data_group.align_on_border(direction=mnm.LEFT, buff=mnm.LARGE_BUFF)))
        self.play(Blink(randy))
        self.play(mnm.Write(ref_group.align_on_border(direction=mnm.LEFT, buff=mnm.LARGE_BUFF)))
        self.play(Blink(morty))
        self.play(mnm.Write(manim_group.align_on_border(direction=mnm.LEFT, buff=mnm.LARGE_BUFF)))
        # self.play(mnm.Write(wechat.align_on_border(direction=mnm.LEFT, buff=mnm.LARGE_BUFF)))
        # wechat1 = mnm.ImageMobject(r"sign_3b1b.png").scale(0.6)
        # wechat2 = mnm.ImageMobject(r"sign_volatility.png").scale(0.6)
        # wechat_group = mnm.Group(wechat1, wechat2).arrange(direction=mnm.RIGHT).align_on_border(mnm.DR)
        # self.play(mnm.FadeIn(wechat_group))


if __name__ == "__main__":
    with mnm.tempconfig(
            {"quality": "production_quality", "preview": True, "frame_rate": 100, "font": "Microsoft YaHei"}):
        print("frame rate:", mnm.config["frame_rate"])
        # scene = AnnotationGroupTest()
        a_market = IndexTimeSeries(A_INDEX_DICT['沪深300'], "2002-07-01", "2022-07-01", fill_flag=True)
        scene = MyMovie(a_market)
        scene.render()
